import os
import torch
import argparse
import numpy as np
import cv2
import torch.utils.data as Data
from utils.binary import assd
from distutils.version import LooseVersion

import sklearn.metrics as metrics
from sklearn.metrics import roc_auc_score

from Models.networks.network import Comprehensive_Atten_Unet


from utils.dataset import coronary_dataset
from utils.transform import coronary_transform

def calculate_Accuracy(confusion):
    confusion=np.asarray(confusion)
    pos = np.sum(confusion, 1).astype(np.float32) # 1 for row
    res = np.sum(confusion, 0).astype(np.float32) # 0 for coloum
    tp = np.diag(confusion).astype(np.float32)
    IU = tp / (pos + res - tp)

    meanIU = np.mean(IU)
    Acc = np.sum(tp) / np.sum(confusion)
    Se = confusion[1][1] / (confusion[1][1]+confusion[1][0])
    Sp = confusion[0][0] / (confusion[0][0]+confusion[0][1])

    return meanIU, Acc, Se, Sp, IU


def test_coronary(test_loader, model):

    Background_IOU = []
    Vessel_IOU = []
    ACC = []
    SE = []
    SP = []
    AUC = []


    model.eval()
    with torch.no_grad():
        for step, (img, lab, image_name) in enumerate(test_loader):

            print("step: ", step)
            image = img.float().cuda()
            # target = lab.float().cuda()

            # output, atten2_map, atten3_map = model(image)  # model output
            # begin_time = time()
            output = model(image)
            # end_time = time()
            # pred_time = end_time - begin_time


            # 计算各种指标
            output = output.cpu().detach().numpy()
            output_01 = np.argmax(output, 1)
            output_save = output_01.transpose(1, 2, 0)
            save_path = "/home/handewei/project/segment/CA-Net/result/pred_filtered/" + image_name[0]
            cv2.imwrite(save_path, output_save * 255)

            output_01 = output_01.reshape([-1])
            target_01 = lab.detach().numpy().astype(np.int64).reshape([-1])
            my_confusion = metrics.confusion_matrix(output_01, target_01).astype(np.float32)
            meanIU, Acc, Se, Sp, IU = calculate_Accuracy(my_confusion)

            Background_IOU.append(IU[0])
            Vessel_IOU.append(IU[1])
            ACC.append(Acc)
            SE.append(Se)
            SP.append(Sp)

            y_pred = output[:, 1, :, :]
            y_pred = y_pred.reshape([-1])

            Auc = roc_auc_score(target_01, y_pred)
            AUC.append(Auc)


            print('Acc: {:.4f} | Se: {:.4f} | Sp: {:.4f} | Auc: {:.4f} | Background_IOU: {:.4f}, vessel_IOU: {:.4f}'
                  .format(Acc, Se, Sp, Auc, IU[0], IU[1]))

    print('avg_ Acc: {:.4f} | Se: {:.4f} | Sp: {:.4f} | Auc: {:.4f} | Background_IOU: {:.4f} |  vessel_IOU: {:.4f} '
          .format(np.average(ACC), np.average(SE), np.average(SP), np.average(AUC),
                  np.average(Background_IOU), np.average(Vessel_IOU)))


if __name__ == '__main__':
    assert LooseVersion(torch.__version__) >= LooseVersion('0.4.0'), 'PyTorch>=0.4.0 is required'

    parser = argparse.ArgumentParser(description='U-net add Attention mechanism for biomedical Dataset')

    # Path related arguments
    parser.add_argument('--root_path', default='./data',
                        help='root directory of data')
    parser.add_argument('--ckpt', default='./saved_models',
                        help='folder to output checkpoints')
    parser.add_argument('--save', default='./result',
                        help='folder to outoput result')
    parser.add_argument('--batch_size', type=int, default=1, metavar='N',
                        help='input batch size for training (default: 16)')
    parser.add_argument('--num_input', default=3, type=int,
                        help='number of input image for each patient')
    parser.add_argument('--num_classes', default=2, type=int,
                        help='number of classes')
    parser.add_argument('--epoch', type=int, default=300, metavar='N',
                        help='choose the specific epoch checkpoints')

    # other arguments
    parser.add_argument('--data', default='coronary', help='choose the dataset')
    parser.add_argument('--out_size', default=(512, 512), help='the output image size')
    parser.add_argument('--att_pos', default='dec', type=str,
                        help='where attention to plug in (enc, dec, enc\&dec)')
    parser.add_argument('--view', default='axial', type=str,
                        help='use what views data to test (for fetal MRI)')


    args = parser.parse_args()
    args.ckpt = os.path.join(args.ckpt, args.data)

    # loading the dataset
    print('loading the {0} dataset ...'.format('test'))
    testset = coronary_dataset(dataset_folder=args.root_path, train_type='test', transform=coronary_transform)
    testloader = Data.DataLoader(dataset=testset, batch_size=args.batch_size, shuffle=False)
    print('Loading is done\n')

    # Define model
    if torch.cuda.is_available():
        print('We can use', torch.cuda.device_count(), 'GPUs to train the network')
        model = Comprehensive_Atten_Unet(args, args.num_input, args.num_classes).cuda()

    # Load the trained best model
    # modelname = './saved_models/coronary/Comp_Atten_Unet/min_loss_coronary_checkpoint.pth.tar'
    modelname = './saved_models/coronary/Comp_Atten_Unet/80_coronary_checkpoint.pth'
    if os.path.isfile(modelname):
        print("=> Loading checkpoint '{}'".format(modelname))
        checkpoint = torch.load(modelname)
        # start_epoch = checkpoint['epoch']

        # multi-GPU transfer to one GPU
        # model_dict = model.state_dict()
        # pretrained_dict = checkpoint['state_dict']
        # from collections import OrderedDict
        # new_state_dict = OrderedDict()
        # for k, v in pretrained_dict.items():
        #     name = k[7:]
        #     new_state_dict[name] = v
        #
        # model_dict.update(new_state_dict)
        # model.load_state_dict(model_dict)

        model.load_state_dict(checkpoint['state_dict'])
        # optimizer.load_state_dict(checkpoint['opt_dict'])
        print("=> Loaded saved the best model at (epoch {})".format(checkpoint['epoch']))
    else:
        print("=> No checkpoint found at '{}'".format(modelname))

    test_coronary(testloader, model)



